import gradio as gr
import torch
import torch.nn.functional as F
from model import LSTMClassifier
from data_loader import text_to_sequence, load_pretrained_embeddings
from utils import load_best_model
import pickle as pkl

# 加载模型和词典
model_path = "best_lstm_model.pth"
word_to_idx_path = "word_to_idx.pth"
embedding_file = "embedding_SougouNews.npz"

# 加载词典
with open("vocab.pkl", "rb") as f:
    word_to_idx = pkl.load(f)
# 加载词嵌入矩阵
embedding_matrix = load_pretrained_embeddings(embedding_file, word_to_idx)
# 初始化模型并加载最佳权重
model = LSTMClassifier(
    vocab_size=len(word_to_idx),
    embed_size=300,
    hidden_size=128,
    output_size=15,
    num_layers=4,
    dropout=0.5,
    embedding_matrix=embedding_matrix,
)
model = load_best_model(model, model_path)


# 定义文本预处理和预测函数
def predict_category(text):
    # 文本预处理
    text = text.strip()
    text_seq = text_to_sequence(
        text, word_to_idx
    )  # 假设text_to_sequence是你定义的函数，将文本转换为序列
    text_tensor = torch.tensor([text_seq], dtype=torch.long)  # 转换为张量

    # 进行预测
    with torch.no_grad():
        output = model(text_tensor)
        probabilities = F.softmax(output, dim=1).cpu().numpy()  # 获取概率
    # 创建结果字典
    # 返回预测结果
    categories = [
        "彩票",
        "食品",
        "财经",
        "汽车",
        "书画",
        "娱乐",
        "科技",
        "国际",
        "健康",
        "旅游",
        "教育",
        "能源",
        "时尚",
        "港澳",
        "台湾",
    ]
    # 类别名称及其概率
    results = {categories[i]: float(prob) for i, prob in enumerate(probabilities[0])}
    # 获取最可能的类别作为文本框输出
    prediction = categories[output.argmax(dim=1).cpu().numpy()[0]]  # 获取预测类别

    return results, str(prediction)  # 返回两个输出值


# 创建Gradio界面
iface = gr.Interface(
    fn=predict_category,
    inputs=gr.Textbox(lines=2, placeholder="输入新闻标题..."),
    outputs=[
        gr.Label(num_top_classes=3, label="预测概率"),
        gr.Textbox(label="预测结果", lines=3),
    ],
    live=False,
    title="新闻标题分类",
    description="输入新闻标题，显示概率前三名的结果",
)

# 启动界面
iface.launch()
